package com.yjzx.llama.tensor;

import com.yjzx.llama.enums.GGMLType;
import com.yjzx.llama.interfaces.AggregateFunction;
import com.yjzx.llama.interfaces.MapFunction;
import com.yjzx.llama.interfaces.MapWithIndexFunction;
import jdk.incubator.vector.FloatVector;
import jdk.incubator.vector.VectorSpecies;

import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;

/**
 * @author yjzx
 * @description FloatTensor
 * @since 2025/5/28
 */
public interface FloatTensor {
    FloatTensor EMPTY = new FloatTensor() {

        @Override
        public int size() {
            return 0;
        }

        @Override
        public float getFloat(int index) {
            return 0;
        }

        @Override
        public void setFloat(int index, float value) {

        }

        @Override
        public FloatVector getFloatVector(VectorSpecies<Float> species, int offset) {
            return null;
        }

        @Override
        public GGMLType type() {
            return null;
        }

        @Override
        public float dot(int thisOffset, FloatTensor that, int thatOffset, int size) {
            return 0;
        }

        @Override
        public void matmul(FloatTensor that, FloatTensor out, int dim0, int dim1) {

        }

        @Override
        public void matmul(int context, FloatTensor[] that, FloatTensor[] out, int dim0, int dim1) {

        }

        @Override
        public float reduce(int thisOffset, int size, float seed, AggregateFunction reduce) {
            return 0;
        }

        @Override
        public float sum(int thisOffset, int size) {
            return 0;
        }

        @Override
        public float max(int thisOffset, int size) {
            return 0;
        }

        @Override
        public void copyTo(int thisOffset, FloatTensor that, int thatOffset, int size) {

        }

        @Override
        public int argmax(int thisOffset, int size) {
            return 0;
        }

        @Override
        public int argmax() {
            return 0;
        }

        @Override
        public FloatTensor mapInPlace(int thisOffset, int size, MapFunction mapFunction) {
            return null;
        }

        @Override
        public FloatTensor mapInPlace(MapFunction mapFunction) {
            return null;
        }

        @Override
        public FloatTensor mapWithIndexInPlace(int thisOffset, int size, MapWithIndexFunction mapWithIndexFunction) {
            return null;
        }

        @Override
        public FloatTensor addInPlace(int thisOffset, FloatTensor that, int thatOffset, int size) {
            return null;
        }

        @Override
        public FloatTensor addInPlace(FloatTensor that) {
            return null;
        }

        @Override
        public FloatTensor multiplyInPlace(int thisOffset, FloatTensor that, int thatOffset, int size) {
            return null;
        }

        @Override
        public FloatTensor multiplyInPlace(FloatTensor that) {
            return null;
        }

        @Override
        public FloatTensor divideInPlace(int thisOffset, int size, float value) {
            return null;
        }

        @Override
        public FloatTensor fillInPlace(int thisOffset, int size, float value) {
            return null;
        }

        @Override
        public FloatTensor softmaxInPlace(int thisOffset, int size) {
            return null;
        }

        @Override
        public FloatTensor saxpyInPlace(int thisOffset, FloatTensor that, int thatOffset, int size, float a) {
            return null;
        }
    };

    /**
     * 读取指定位置的short
     *
     * @param memorySegment 内存地址
     * @param offset        偏移量
     * @return
     */
    default short readShort(MemorySegment memorySegment, long offset) {
        // The MemorySegment.get* methods should be used instead.
        return memorySegment.get(ValueLayout.JAVA_SHORT, offset);
    }

    /**
     * 读取指定位置的byte
     *
     * @param memorySegment 内存地址
     * @param offset        偏移量
     * @return
     */
    default byte readByte(MemorySegment memorySegment, long offset) {
        // The MemorySegment.get* methods should be used instead.
        return memorySegment.get(ValueLayout.JAVA_BYTE, offset);
    }
    /**
     * 长度
     *
     * @return
     */
    int size();

    /**
     * 获取元素
     *
     * @param index
     * @return
     */
    float getFloat(int index);

    /**
     * 修改元素
     *
     * @param index
     * @param value
     */
    void setFloat(int index, float value);

    /**
     * 根据向量种类对数组进行填充
     *
     * @param species 向量类型
     * @param offset  偏移量位置开始
     * @return 填充后的向量值
     */
    FloatVector getFloatVector(VectorSpecies<Float> species, int offset);

    /**
     * 模型类型
     *
     * @return
     */
    GGMLType type();
    /**
     * 点积函数
     *
     * @param thisOffset 张量1偏移起始量
     * @param that       张量2
     * @param thatOffset 张量2偏移起始量
     * @param size       数量
     * @return 对俩张量偏移后的数值进行求积后求和相加
     */

    float dot(int thisOffset, FloatTensor that, int thatOffset, int size);
    /**
     * 矩阵乘积
     *
     * @param that 参数1
     * @param out  输出参数
     * @param dim0 0-dim0的每个值对
     * @param dim1
     */
    void matmul(FloatTensor that, FloatTensor out, int dim0, int dim1);

    /**
     * 双重矩阵乘积
     *
     * @param context 0-dim0*context
     * @param that    参数1
     * @param out     输出参数
     * @param dim0    第一个
     * @param dim1    第二个
     */
    void matmul(int context, FloatTensor[] that, FloatTensor[] out, int dim0, int dim1);

    /**
     * 对
     *
     * @param thisOffset 当前对象偏移量
     * @param size       长度
     * @param seed       运算参数
     * @param reduce     进行的运算
     * @return 对当前对象指定的索引位置的值与seed在值size长度遍历进行reduce运算
     */
    float reduce(int thisOffset, int size, float seed, AggregateFunction reduce);


    float sum(int thisOffset, int size);

    float max(int thisOffset, int size);

    /**
     * 对that的到that Offset到size长度偏移量的值替换为 this的偏移量及size长度偏移量的值
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     */
    void copyTo(int thisOffset, FloatTensor that, int thatOffset, int size);

    /**
     * @param thisOffset 偏移量
     * @param size       长度
     * @return 获取thisOffset到指定size长度中最大索引值
     */
    int argmax(int thisOffset, int size);

    /**
     * @return 获取有最大值的索引（最后一个）
     */
    int argmax();

    FloatTensor mapInPlace(int thisOffset, int size, MapFunction mapFunction);

    FloatTensor mapInPlace(MapFunction mapFunction);

    FloatTensor mapWithIndexInPlace(int thisOffset, int size, MapWithIndexFunction mapWithIndexFunction);

    /**
     * 对向量进行求和
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     * @return
     */
    FloatTensor addInPlace(int thisOffset, FloatTensor that, int thatOffset, int size);

    /**
     * 对向量进行求和
     *
     * @param that
     * @return
     */
    FloatTensor addInPlace(FloatTensor that);

    /**
     * 对向量进行乘积
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     * @return
     */
    FloatTensor multiplyInPlace(int thisOffset, FloatTensor that, int thatOffset, int size);

    /**
     * 对向量进行乘积
     *
     * @param that
     * @return
     */
    FloatTensor multiplyInPlace(FloatTensor that);
    /**
     * 对向量进行求商
     *
     * @param thisOffset
     * @param size
     * @param value
     * @return
     */
    FloatTensor divideInPlace(int thisOffset, int size, float value);

    /**
     * 填充所有元素为value
     *
     * @param thisOffset 开始位置
     * @param size       填充数量
     * @param value      填充值
     * @return
     */
    FloatTensor fillInPlace(int thisOffset, int size, float value) ;

    FloatTensor softmaxInPlace(int thisOffset, int size);

    /**
     * 对向量进行a倍乘积+that值
     *
     * @param thisOffset
     * @param that
     * @param thatOffset
     * @param size
     * @param a
     * @return
     */
    FloatTensor saxpyInPlace(int thisOffset, FloatTensor that, int thatOffset, int size, float a);

}
